import pandas as pd
from datasets import Dataset, DatasetDict


def encode_label():
    """
    把样本中的标签进行编码
    把数据写入到dataset格式中
    :return:
    """

    label_to_index = {"O": 0, "B-dis": 1, "I-dis": 2, "B-sym": 3, "I-sym": 4}
    
    # 加载csv文件
    train_data = pd.read_csv('doctor_offline/ner_model/ner_data/train.csv')
    valid_data = pd.read_csv('doctor_offline/ner_model/ner_data/valid.csv')
    train_data = Dataset.from_pandas(train_data)
    valid_data = Dataset.from_pandas(valid_data)
    corpus_data = DatasetDict({'train': train_data, 'valid': valid_data})
    
    # 处理标签
    def data_handler(data_labels, data_inputs):
        # 把标签由字符串 映射成数字
        data_label_ids = []
        for labels in data_labels:
            label_ids = []
            for label in labels.split():
                label_ids.append(label_to_index[label])
            data_label_ids.append(label_ids)

        return {'data_labels': data_label_ids, 'data_inputs': data_inputs}
    
    
    corpus_data = corpus_data.map(data_handler, input_columns=['data_labels', 'data_inputs'], batched=True)
	# 注意，这里map执行完成后，需要用corpus_data变量接收，否则存入的数据仍然是原始数据
    corpus_data.save_to_disk('doctor_offline/ner_model/ner_data/bilstm_crf_data')
                

        
        
if __name__ == '__main__':
    encode_label()






